{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8f8e89ca-f181-4dd2-bdb6-84126ccc881d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import DataLoader\n",
    "import torchvision.datasets as dsets\n",
    "import torchvision.transforms as transforms\n",
    "batch_size = 100\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ee994f48-7686-4364-8f97-f21a15a1eccc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_data: torch.Size([60000, 28, 28])\n",
      "train_labels: torch.Size([60000])\n",
      "test_data: torch.Size([10000, 28, 28])\n",
      "test_labels: torch.Size([10000])\n",
      "批次的尺寸: 100\n",
      "load_train_data: torch.Size([60000, 28, 28])\n",
      "load_train_labels: torch.Size([60000])\n"
     ]
    }
   ],
   "source": [
    "#下载MNIST dataset\n",
    "train_dataset = dsets.MNIST(root='/data/project/python/torch/data/mnist',\n",
    "                            train=True,\n",
    "                            transform=transforms.ToTensor(),\n",
    "                            download=True)\n",
    "test_dataset = dsets.MNIST(root='/data/project/python/torch/data/mnist',\n",
    "                            train=False,\n",
    "                            transform=transforms.ToTensor(),\n",
    "                            download=True)\n",
    "\n",
    "#加载数据\n",
    "train_loader = torch.utils.data.DataLoader(dataset=train_dataset,\n",
    "                                           batch_size=batch_size,\n",
    "                                           shuffle=True)\n",
    "test_loader = torch.utils.data.DataLoader(dataset=test_dataset,\n",
    "                                           batch_size=batch_size,\n",
    "                                           shuffle=True)\n",
    "\n",
    "#原始数据\n",
    "print(\"train_data:\", train_dataset.data.size())            #train_data\n",
    "print(\"train_labels:\", train_dataset.targets.size())       #train_labels\n",
    "print(\"test_data:\", test_dataset.data.size())              #test_data\n",
    "print(\"test_labels:\", test_dataset.targets.size())         #test_labels\n",
    "#数据打乱取小批次\n",
    "print('批次的尺寸:', train_loader.batch_size)\n",
    "print('load_train_data:', train_loader.dataset.data.shape)\n",
    "print('load_train_labels:', train_loader.dataset.targets.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "36ded824-9deb-4e01-b3c3-8948bff06fb8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 定义神经网络\n",
    "import torch.nn as nn\n",
    "import torch\n",
    "\n",
    "input_size = 784    #mnist的像素为28*28\n",
    "hidden_size = 500\n",
    "num_classes = 10   #输入10个类别\n",
    "\n",
    "#创建神经网络模型\n",
    "class Neural_net(nn.Module):\n",
    "    def __init__(self, input_num, hidden_size, out_put):\n",
    "        super(Neural_net, self).__init__()\n",
    "        self.layer1 = nn.Linear(input_num, hidden_size)\n",
    "        self.layer2 = nn.Linear(hidden_size, out_put)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.layer1(x)      #输入层到隐藏层的线性计算\n",
    "        out = torch.relu(out)     #隐藏层激活\n",
    "        out = self.layer2(out)    #输出层，注意，输出层直接loss\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7c03eea7-e7ae-43b3-83a6-ec4a798e8089",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Neural_net(\n",
      "  (layer1): Linear(in_features=784, out_features=500, bias=True)\n",
      "  (layer2): Linear(in_features=500, out_features=10, bias=True)\n",
      ")\n",
      "current epoch = 0\n",
      "current loss = 2.31438\n",
      "current loss = 0.56605\n",
      "current loss = 0.40846\n",
      "current loss = 0.36222\n",
      "current loss = 0.42443\n",
      "current loss = 0.26781\n",
      "current epoch = 1\n",
      "current loss = 0.41010\n",
      "current loss = 0.32272\n",
      "current loss = 0.22890\n",
      "current loss = 0.20686\n",
      "current loss = 0.33159\n",
      "current loss = 0.16995\n",
      "current epoch = 2\n",
      "current loss = 0.19951\n",
      "current loss = 0.13445\n",
      "current loss = 0.18122\n",
      "current loss = 0.11838\n",
      "current loss = 0.16861\n",
      "current loss = 0.12138\n",
      "current epoch = 3\n",
      "current loss = 0.11336\n",
      "current loss = 0.12318\n",
      "current loss = 0.21791\n",
      "current loss = 0.13869\n",
      "current loss = 0.17391\n",
      "current loss = 0.21060\n",
      "current epoch = 4\n",
      "current loss = 0.24302\n",
      "current loss = 0.16223\n",
      "current loss = 0.15516\n",
      "current loss = 0.17002\n",
      "current loss = 0.09116\n",
      "current loss = 0.20314\n",
      "finished training\n"
     ]
    }
   ],
   "source": [
    "# 训练\n",
    "from torch.autograd import Variable\n",
    "import numpy as np\n",
    "\n",
    "net = Neural_net(input_size, hidden_size, num_classes)\n",
    "print(net)\n",
    "\n",
    "learning_rate = 1e-1\n",
    "num_epoches = 5\n",
    "\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.SGD(net.parameters(), lr=learning_rate)\n",
    "for epoch in range(num_epoches):\n",
    "    print('current epoch = %d' % epoch)\n",
    "    for i, (images, labels) in enumerate(train_loader):\n",
    "        images = Variable(images.view(-1, 28*28))\n",
    "        lables = Variable(labels)\n",
    "\n",
    "        outputs = net(images)\n",
    "        loss = criterion(outputs, labels)\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        if i % 100 == 0:\n",
    "            print('current loss = %.5f' % loss.item())\n",
    "            \n",
    "print('finished training')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7bbc8df-1902-47eb-9b6d-7328cb186e8a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py-3.12_torch_24",
   "language": "python",
   "name": "py3.12_torch24"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
